
# Author: qiang
import time

import socket
class WindowRateTrack():
  def __init__(self, window_size: int) -> None:
    self._window_size = window_size
    self._rate_limit_ = None
    self._sum = 0
    self._current_window_index = 0
    self._offset = 0

  def set_rate_limit_info(self, rate: int, rate_limit_remaining: int):
    self._rate_limit_ = rate
    self._sum = rate - rate_limit_remaining

  def update_time(self, time: int):
    window_index = int((time - self._offset) / self._window_size)
    if window_index > self._current_window_index:
      self._sum = 0
      self._current_window_index = window_index

  def max_req(self, ):
    return self._rate_limit_

  def push(self, num_reqs: int, timestamp: int):
    self.update_time(timestamp)
    self._sum += num_reqs

  def get_remaining(self, time: int):
    self.update_time(time)
    return max(0, self._rate_limit_ - self._sum)


import threading
class IpRateTrackManager():
  def __init__(self, name) -> None:
    self._name = name
    self._ip_map = {}
    self._lock = threading.Lock()

  def exchange_name(self) -> str:
    return self._name

  def add_ip(self, ip: str, window_size: int, rate: int, rate_limit_remaining: int) -> bool:
    self._lock.acquire()

    if ip not in self._ip_map:
      self._ip_map[ip] = WindowRateTrack(window_size)
      self._ip_map[ip].set_rate_limit_info(rate, rate_limit_remaining)

    self._lock.release()

  def del_ip(self, ip):
    self._lock.acquire()
    del (self[ip])
    self._lock.release()

  def get_ip(self, num_resq: int) -> str:
    self._lock.acquire()
    rate_track = None
    for ip, track in self._ip_map.items():
      assert (isinstance(track, WindowRateTrack))
      remaining = track.get_remaining(time.time_ns())
      if remaining - num_resq > 0:
        rate_track = track
        break

    if rate_track is None:
      ip = ""
    else:
      rate_track.push(num_resq, time.time_ns())
    self._lock.release()
    return ip
  
  def ip_size(self):
    return len(self._ip_map)


def get_host_ips():
  addrs = socket.getaddrinfo(socket.gethostname(), None, family=socket.AF_INET)
  ips = set()
  for item in addrs:
    ips.add(item[4][0])

  return list(ips)

if __name__ == "__main__":
  import time
  kOneSecToNs = 1000000000
  # print(time.time_ns())
  # track = WindowRateTrack(kOneSecToNs)
  # track.set_rate_limit_info(10, 10)
  # while True:
  #   num = track.get_remaining(time.time_ns())
  #   print(num)
  #   if num == 0:
  #     print("limit error")
  #   track.push(1, time.time_ns())
  #   time.sleep(0.09)

  ip_m = IpRateTrackManager("s")
  ip_m.add_ip("192.168.120.1", kOneSecToNs, 1200, 1200)
  ip_m.add_ip("192.168.120.2", kOneSecToNs, 1200, 1200)

  i = 0
  timestamp = time.time_ns()
  index = int(timestamp / kOneSecToNs)
  while True:
    nums_req = 0
    if i % 3 == 0:
      nums_req = 40
    else:
      nums_req = 3

    ip = ip_m.get_ip(nums_req)
    if len(ip) == 0:
      print("api limit err")
    else:
      print(ip)

    time.sleep(0.01)
    timestamp = time.time_ns()
    index1 = int(timestamp / kOneSecToNs)
    if index1 != index:
      index = index1
